#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MCP不确定性决策引擎
实现蒙特卡洛模拟、鲁棒优化、风险分析等不确定性决策功能
"""

import numpy as np
import random
import time
from typing import Dict, List, Any, Optional, Callable, Tuple
from scipy import stats, optimize
from scipy.stats import norm, uniform, triang, expon
import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed

# 尝试相对导入，如果失败则使用绝对导入
try:
    from .models import (
        OptimizationProblem, UncertaintyParameter, UncertaintyType,
        RiskProfile, ScenarioAnalysis, MonteCarloResult, 
        RobustSolution, UncertaintyResult, DecisionVariable
    )
except ImportError:
    from models import (
        OptimizationProblem, UncertaintyParameter, UncertaintyType,
        RiskProfile, ScenarioAnalysis, MonteCarloResult, 
        RobustSolution, UncertaintyResult, DecisionVariable
    )


class UncertaintyEngine:
    """不确定性决策引擎"""
    
    def __init__(self, random_seed: Optional[int] = None):
        """初始化不确定性引擎"""
        self.random_seed = random_seed
        if random_seed is not None:
            np.random.seed(random_seed)
            random.seed(random_seed)
        
        # 支持的概率分布
        self.distributions = {
            UncertaintyType.NORMAL: self._sample_normal,
            UncertaintyType.UNIFORM: self._sample_uniform,
            UncertaintyType.TRIANGULAR: self._sample_triangular,
            UncertaintyType.EXPONENTIAL: self._sample_exponential
        }
    
    def monte_carlo_simulation(self, problem: OptimizationProblem, 
                             objective_function: Callable,
                             simulation_count: int = 10000,
                             confidence_levels: List[float] = [0.95, 0.99],
                             parallel: bool = True) -> MonteCarloResult:
        """执行蒙特卡洛模拟"""
        start_time = time.time()
        
        # 生成样本
        samples = self._generate_samples(problem, simulation_count)
        
        # 计算目标函数值
        if parallel and simulation_count > 1000:
            objective_samples = self._evaluate_parallel(objective_function, samples)
        else:
            objective_samples = [objective_function(sample) for sample in samples]
        
        # 计算风险概况
        risk_profile = self._calculate_risk_profile(objective_samples, confidence_levels)
        
        # 收敛性分析
        convergence_analysis = self._analyze_convergence(objective_samples)
        
        execution_time = time.time() - start_time
        
        return MonteCarloResult(
            simulation_count=simulation_count,
            objective_samples=objective_samples,
            solution_samples=samples,
            risk_profile=risk_profile,
            convergence_analysis=convergence_analysis,
            execution_time=execution_time,
            random_seed=self.random_seed
        )
    
    def robust_optimization(self, problem: OptimizationProblem,
                          objective_function: Callable,
                          scenario_count: int = 100,
                          robustness_measure: str = "worst_case") -> RobustSolution:
        """执行鲁棒优化"""
        start_time = time.time()
        
        # 生成不确定性情景
        scenarios = self._generate_scenarios(problem, scenario_count)
        
        # 定义鲁棒目标函数
        if robustness_measure == "worst_case":
            robust_objective = self._worst_case_objective
        elif robustness_measure == "expected_worst":
            robust_objective = self._expected_worst_objective
        else:
            robust_objective = self._minimax_regret_objective
        
        # 求解鲁棒优化问题
        robust_solution = self._solve_robust_problem(
            problem, objective_function, scenarios, robust_objective
        )
        
        # 评估鲁棒解在各情景下的表现
        scenario_analyses = self._evaluate_scenarios(
            robust_solution, scenarios, objective_function
        )
        
        # 敏感性分析
        sensitivity_analysis = self._sensitivity_analysis(
            problem, robust_solution, objective_function
        )
        
        execution_time = time.time() - start_time
        
        # 计算性能指标
        objective_values = [obj_val for _, obj_val in 
                          [(s.solution, objective_function(s.solution)) for s in scenario_analyses]]
        
        return RobustSolution(
            solution=robust_solution,
            worst_case_objective=min(objective_values),
            best_case_objective=max(objective_values),
            expected_objective=np.mean(objective_values),
            robustness_measure=np.std(objective_values),
            scenario_performance=scenario_analyses,
            uncertainty_budget=self._calculate_uncertainty_budget(problem),
            sensitivity_analysis=sensitivity_analysis
        )
    
    def scenario_analysis(self, problem: OptimizationProblem,
                         objective_function: Callable,
                         scenarios: List[Dict[str, Any]]) -> List[ScenarioAnalysis]:
        """执行情景分析"""
        analyses = []
        
        for i, scenario_params in enumerate(scenarios):
            scenario_name = scenario_params.get('name', f'Scenario_{i+1}')
            probability = scenario_params.get('probability', 1.0 / len(scenarios))
            
            # 在该情景下求解优化问题
            scenario_solution = self._solve_scenario_problem(
                problem, objective_function, scenario_params
            )
            
            # 计算目标函数值
            objective_values = {
                'primary': objective_function(scenario_solution)
            }
            
            # 计算风险指标
            risk_metrics = self._calculate_scenario_risk(scenario_solution, scenario_params)
            
            analyses.append(ScenarioAnalysis(
                scenario_name=scenario_name,
                probability=probability,
                objective_values=objective_values,
                solution=scenario_solution,
                risk_metrics=risk_metrics,
                description=scenario_params.get('description', '')
            ))
        
        return analyses
    
    def uncertainty_impact_analysis(self, problem: OptimizationProblem,
                                   objective_function: Callable,
                                   base_solution: Dict[str, float]) -> Dict[str, float]:
        """分析不确定性对目标函数的影响"""
        impact_analysis = {}
        base_objective = objective_function(base_solution)
        
        for uncertainty_param in problem.uncertainty_params:
            # 计算该不确定性参数的影响
            param_impact = self._calculate_parameter_impact(
                uncertainty_param, objective_function, base_solution, base_objective
            )
            impact_analysis[uncertainty_param.name] = param_impact
        
        return impact_analysis
    
    def generate_decision_recommendations(self, uncertainty_result: UncertaintyResult) -> List[str]:
        """生成决策建议"""
        recommendations = []
        
        # 基于风险等级的建议
        risk_level = uncertainty_result.get_overall_risk_level()
        if risk_level == "高风险":
            recommendations.append("建议采用保守策略，优先考虑风险控制")
            recommendations.append("建议增加不确定性缓冲，提高方案鲁棒性")
        elif risk_level == "中等风险":
            recommendations.append("建议平衡风险与收益，采用适度进取策略")
            recommendations.append("建议建立风险监控机制，及时调整策略")
        else:
            recommendations.append("风险较低，可以采用相对进取的策略")
        
        # 基于蒙特卡洛结果的建议
        if uncertainty_result.monte_carlo_result:
            mc_result = uncertainty_result.monte_carlo_result
            cv = mc_result.risk_profile.standard_deviation / mc_result.risk_profile.expected_value
            if cv > 0.5:
                recommendations.append("目标函数变异系数较高，建议重新评估决策变量范围")
        
        # 基于鲁棒解的建议
        if uncertainty_result.robust_solution:
            robust_score = uncertainty_result.robust_solution.get_robustness_score()
            if robust_score < 60:
                recommendations.append("当前解的鲁棒性较差，建议考虑更保守的策略")
            elif robust_score > 90:
                recommendations.append("当前解具有很好的鲁棒性，可以考虑适度优化收益")
        
        # 基于情景分析的建议
        if uncertainty_result.scenario_analyses:
            high_prob_scenarios = [s for s in uncertainty_result.scenario_analyses if s.probability > 0.3]
            if high_prob_scenarios:
                recommendations.append(f"重点关注高概率情景：{', '.join([s.scenario_name for s in high_prob_scenarios])}")
        
        return recommendations
    
    # ========================================================================
    # 私有方法
    # ========================================================================
    
    def _generate_samples(self, problem: OptimizationProblem, count: int) -> List[Dict[str, float]]:
        """生成随机样本"""
        samples = []
        
        for _ in range(count):
            sample = {}
            
            # 为每个决策变量生成样本值
            for var in problem.variables:
                if var.variable_type.value == "continuous":
                    sample[var.name] = np.random.uniform(
                        var.lower_bound or 0, var.upper_bound or 1
                    )
                elif var.variable_type.value == "integer":
                    sample[var.name] = np.random.randint(
                        int(var.lower_bound or 0), int(var.upper_bound or 10) + 1
                    )
                elif var.variable_type.value == "binary":
                    sample[var.name] = np.random.choice([0, 1])
            
            # 为不确定性参数生成样本值
            for uncertainty_param in problem.uncertainty_params:
                sample[uncertainty_param.name] = self._sample_uncertainty_parameter(uncertainty_param)
            
            samples.append(sample)
        
        return samples
    
    def _sample_uncertainty_parameter(self, param: UncertaintyParameter) -> float:
        """为不确定性参数生成样本"""
        if param.uncertainty_type in self.distributions:
            return self.distributions[param.uncertainty_type](param.parameters)
        else:
            raise ValueError(f"不支持的不确定性类型: {param.uncertainty_type}")
    
    def _sample_distribution(self, param: UncertaintyParameter, size: int) -> np.ndarray:
        """为不确定性参数生成多个样本"""
        if param.uncertainty_type == UncertaintyType.NORMAL:
            mean = param.parameters.get('mean', 0.0)
            std = param.parameters.get('std', 1.0)
            return np.random.normal(mean, std, size)
        
        elif param.uncertainty_type == UncertaintyType.UNIFORM:
            min_val = param.parameters.get('min', 0.0)
            max_val = param.parameters.get('max', 1.0)
            return np.random.uniform(min_val, max_val, size)
        
        elif param.uncertainty_type == UncertaintyType.TRIANGULAR:
            min_val = param.parameters.get('min', 0.0)
            mode = param.parameters.get('mode', 0.5)
            max_val = param.parameters.get('max', 1.0)
            return np.random.triangular(min_val, mode, max_val, size)
        
        elif param.uncertainty_type == UncertaintyType.EXPONENTIAL:
            rate = param.parameters.get('rate', 1.0)
            return np.random.exponential(1.0/rate, size)
        
        else:
            raise ValueError(f"不支持的不确定性类型: {param.uncertainty_type}")
    
    def _sample_normal(self, params: Dict[str, float]) -> float:
        """正态分布采样"""
        return np.random.normal(params['mean'], params['std'])
    
    def _sample_uniform(self, params: Dict[str, float]) -> float:
        """均匀分布采样"""
        return np.random.uniform(params['min'], params['max'])
    
    def _sample_triangular(self, params: Dict[str, float]) -> float:
        """三角分布采样"""
        return np.random.triangular(params['min'], params['mode'], params['max'])
    
    def _sample_exponential(self, params: Dict[str, float]) -> float:
        """指数分布采样"""
        return np.random.exponential(1.0 / params['rate'])
    
    def _evaluate_parallel(self, objective_function: Callable, samples: List[Dict[str, float]]) -> List[float]:
        """并行计算目标函数值"""
        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = [executor.submit(objective_function, sample) for sample in samples]
            results = [future.result() for future in as_completed(futures)]
        return results
    
    def _calculate_risk_profile(self, samples: List[float], confidence_levels: List[float]) -> RiskProfile:
        """计算风险概况"""
        samples_array = np.array(samples)
        
        # 基本统计量
        expected_value = np.mean(samples_array)
        std_dev = np.std(samples_array)
        skewness = stats.skew(samples_array)
        kurtosis = stats.kurtosis(samples_array)
        
        # VaR和CVaR计算
        var_95 = np.percentile(samples_array, 5)  # 5%分位数
        var_99 = np.percentile(samples_array, 1)  # 1%分位数
        
        # CVaR (条件风险价值)
        cvar_95 = np.mean(samples_array[samples_array <= var_95])
        cvar_99 = np.mean(samples_array[samples_array <= var_99])
        
        # 置信区间
        confidence_intervals = {}
        for level in confidence_levels:
            alpha = 1 - level
            lower = np.percentile(samples_array, 100 * alpha / 2)
            upper = np.percentile(samples_array, 100 * (1 - alpha / 2))
            confidence_intervals[f"{int(level*100)}%"] = (lower, upper)
        
        return RiskProfile(
            var_95=var_95,
            var_99=var_99,
            cvar_95=cvar_95,
            cvar_99=cvar_99,
            expected_value=expected_value,
            standard_deviation=std_dev,
            skewness=skewness,
            kurtosis=kurtosis,
            confidence_intervals=confidence_intervals
        )
    
    def _analyze_convergence(self, samples: List[float]) -> Dict[str, Any]:
        """分析收敛性"""
        n = len(samples)
        running_means = [np.mean(samples[:i+1]) for i in range(n)]
        
        # 计算收敛指标
        final_mean = running_means[-1]
        convergence_error = [abs(mean - final_mean) / abs(final_mean) if final_mean != 0 else 0 
                           for mean in running_means]
        
        # 找到收敛点（误差小于1%）
        convergence_point = None
        for i, error in enumerate(convergence_error):
            if error < 0.01 and i > n // 10:  # 至少10%的样本后才考虑收敛
                convergence_point = i
                break
        
        return {
            "running_means": running_means[-100:],  # 只保留最后100个点
            "convergence_errors": convergence_error[-100:],
            "convergence_point": convergence_point,
            "final_standard_error": np.std(samples) / np.sqrt(n),
            "is_converged": convergence_point is not None
        }
    
    def _generate_scenarios(self, problem: OptimizationProblem, count: int) -> List[Dict[str, Any]]:
        """生成不确定性情景"""
        scenarios = []
        
        for i in range(count):
            scenario = {"scenario_id": i}
            
            for uncertainty_param in problem.uncertainty_params:
                # 为每个不确定性参数生成情景值
                scenario[uncertainty_param.name] = self._sample_uncertainty_parameter(uncertainty_param)
            
            scenarios.append(scenario)
        
        return scenarios
    
    def _worst_case_objective(self, solution: Dict[str, float], scenarios: List[Dict[str, Any]], 
                            objective_function: Callable) -> float:
        """最坏情况目标函数"""
        worst_value = float('inf')
        
        for scenario in scenarios:
            combined_params = {**solution, **scenario}
            value = objective_function(combined_params)
            worst_value = min(worst_value, value)
        
        return worst_value
    
    def _expected_worst_objective(self, solution: Dict[str, float], scenarios: List[Dict[str, Any]], 
                                objective_function: Callable) -> float:
        """期望最坏情况目标函数"""
        values = []
        
        for scenario in scenarios:
            combined_params = {**solution, **scenario}
            value = objective_function(combined_params)
            values.append(value)
        
        # 返回最坏10%情况的期望值
        sorted_values = sorted(values)
        worst_10_percent = sorted_values[:max(1, len(sorted_values) // 10)]
        return np.mean(worst_10_percent)
    
    def _minimax_regret_objective(self, solution: Dict[str, float], scenarios: List[Dict[str, Any]], 
                                objective_function: Callable) -> float:
        """最小最大后悔目标函数"""
        max_regret = 0
        
        for scenario in scenarios:
            # 计算在该情景下的最优值
            optimal_value = self._solve_scenario_optimal(scenario, objective_function)
            
            # 计算当前解在该情景下的值
            combined_params = {**solution, **scenario}
            current_value = objective_function(combined_params)
            
            # 计算后悔值
            regret = optimal_value - current_value
            max_regret = max(max_regret, regret)
        
        return max_regret
    
    def _solve_robust_problem(self, problem: OptimizationProblem, objective_function: Callable,
                            scenarios: List[Dict[str, Any]], robust_objective: Callable) -> Dict[str, float]:
        """求解鲁棒优化问题"""
        # 简化实现：使用随机搜索
        best_solution = None
        best_objective = float('-inf')
        
        for _ in range(1000):  # 随机搜索1000次
            # 生成随机解
            solution = {}
            for var in problem.variables:
                if var.variable_type.value == "continuous":
                    solution[var.name] = np.random.uniform(
                        var.lower_bound or 0, var.upper_bound or 1
                    )
                elif var.variable_type.value == "integer":
                    solution[var.name] = np.random.randint(
                        int(var.lower_bound or 0), int(var.upper_bound or 10) + 1
                    )
                elif var.variable_type.value == "binary":
                    solution[var.name] = np.random.choice([0, 1])
            
            # 评估鲁棒目标函数
            objective_value = robust_objective(solution, scenarios, objective_function)
            
            if objective_value > best_objective:
                best_objective = objective_value
                best_solution = solution.copy()
        
        return best_solution
    
    def _evaluate_scenarios(self, solution: Dict[str, float], scenarios: List[Dict[str, Any]],
                          objective_function: Callable) -> List[ScenarioAnalysis]:
        """评估解在各情景下的表现"""
        analyses = []
        
        for i, scenario in enumerate(scenarios):
            combined_params = {**solution, **scenario}
            objective_value = objective_function(combined_params)
            
            analyses.append(ScenarioAnalysis(
                scenario_name=f"Scenario_{i+1}",
                probability=1.0 / len(scenarios),
                objective_values={"primary": objective_value},
                solution=solution,
                risk_metrics={"objective_value": objective_value}
            ))
        
        return analyses
    
    def _sensitivity_analysis(self, problem: OptimizationProblem, solution: Dict[str, float],
                            objective_function: Callable) -> Dict[str, float]:
        """敏感性分析"""
        sensitivity = {}
        base_value = objective_function(solution)
        
        for var_name, var_value in solution.items():
            # 计算对该变量的敏感性
            perturbed_solution = solution.copy()
            perturbation = 0.01 * abs(var_value) if var_value != 0 else 0.01
            
            perturbed_solution[var_name] = var_value + perturbation
            perturbed_value = objective_function(perturbed_solution)
            
            sensitivity[var_name] = (perturbed_value - base_value) / perturbation
        
        return sensitivity
    
    def _calculate_uncertainty_budget(self, problem: OptimizationProblem) -> Dict[str, float]:
        """计算不确定性预算"""
        budget = {}
        
        for uncertainty_param in problem.uncertainty_params:
            if uncertainty_param.uncertainty_type == UncertaintyType.NORMAL:
                # 对于正态分布，使用3σ作为预算
                budget[uncertainty_param.name] = 3 * uncertainty_param.parameters.get('std', 1.0)
            elif uncertainty_param.uncertainty_type == UncertaintyType.UNIFORM:
                # 对于均匀分布，使用范围的一半作为预算
                range_val = uncertainty_param.parameters.get('max', 1) - uncertainty_param.parameters.get('min', 0)
                budget[uncertainty_param.name] = range_val / 2
            else:
                budget[uncertainty_param.name] = 1.0  # 默认预算
        
        return budget
    
    def _solve_scenario_problem(self, problem: OptimizationProblem, objective_function: Callable,
                              scenario_params: Dict[str, Any]) -> Dict[str, float]:
        """在特定情景下求解优化问题"""
        # 简化实现：随机搜索
        best_solution = None
        best_objective = float('-inf')
        
        for _ in range(100):
            solution = {}
            for var in problem.variables:
                if var.variable_type.value == "continuous":
                    solution[var.name] = np.random.uniform(
                        var.lower_bound or 0, var.upper_bound or 1
                    )
                elif var.variable_type.value == "integer":
                    solution[var.name] = np.random.randint(
                        int(var.lower_bound or 0), int(var.upper_bound or 10) + 1
                    )
                elif var.variable_type.value == "binary":
                    solution[var.name] = np.random.choice([0, 1])
            
            # 添加情景参数
            combined_params = {**solution, **scenario_params}
            objective_value = objective_function(combined_params)
            
            if objective_value > best_objective:
                best_objective = objective_value
                best_solution = solution.copy()
        
        return best_solution
    
    def _calculate_scenario_risk(self, solution: Dict[str, float], scenario_params: Dict[str, Any]) -> Dict[str, float]:
        """计算情景风险指标"""
        # 简化的风险指标计算
        risk_metrics = {
            "volatility": np.std(list(solution.values())),
            "max_exposure": max(solution.values()) if solution.values() else 0,
            "diversification": len([v for v in solution.values() if v > 0]) / len(solution) if solution else 0
        }
        
        return risk_metrics
    
    def _calculate_parameter_impact(self, uncertainty_param: UncertaintyParameter,
                                  objective_function: Callable, base_solution: Dict[str, float],
                                  base_objective: float) -> float:
        """计算不确定性参数的影响"""
        # 生成该参数的样本
        samples = [self._sample_uncertainty_parameter(uncertainty_param) for _ in range(100)]
        
        # 计算目标函数值的变化
        objective_changes = []
        for sample_value in samples:
            modified_solution = base_solution.copy()
            modified_solution[uncertainty_param.name] = sample_value
            new_objective = objective_function(modified_solution)
            objective_changes.append(abs(new_objective - base_objective))
        
        # 返回平均影响
        return np.mean(objective_changes)
    
    def _solve_scenario_optimal(self, scenario: Dict[str, Any], objective_function: Callable) -> float:
        """求解特定情景下的最优值（简化实现）"""
        # 这里应该使用专门的优化算法，简化为随机搜索
        best_value = float('-inf')
        
        for _ in range(50):
            # 生成随机解（这里需要根据具体问题调整）
            random_solution = {f"x_{i}": np.random.uniform(0, 1) for i in range(5)}
            combined_params = {**random_solution, **scenario}
            value = objective_function(combined_params)
            best_value = max(best_value, value)
        
        return best_value